# 构建⼀个图。
# 图中两个步骤：
# 第⼀步让⼤模型推荐⼀个有名的作家，
# 第⼆步，让⼤模型⽤推荐的作家的⻛格写⼀个100字以内的笑话。
from typing import TypedDict

from langchain_openai import ChatOpenAI
from typing_extensions import NotRequired
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.constants import START, END
from langgraph.graph import StateGraph

# 初始化大语言模型客户端
# model: 使用的模型名称
# base_url: API基础URL地址
# openai_api_key: API密钥
llm = ChatOpenAI(
    model="deepseek-chat",
    base_url="https://api.deepseek.com/v1",
    openai_api_key="sk-ddad02a64c4c47ea9f6f526ef47cb602",
)


class State(TypedDict):
    """
    定义图中传递的状态数据结构

    Attributes:
        author (NotRequired[str]): 作家名称，可选字段
        joke (NotRequired[str]): 笑话内容，可选字段
    """
    author: NotRequired[str]
    joke: NotRequired[str]


def author_node(state: State):
    """
    推荐作家的节点函数

    Args:
        state (State): 当前状态对象

    Returns:
        dict: 包含推荐作家名称的字典
    """
    # 构造提示词，请求推荐一位受欢迎的作家
    prompt = "帮我推荐⼀位受⼈们欢迎的作家。只需要给出作家的名字即可。"
    # 调用大语言模型获取作家推荐
    author = llm.invoke(prompt)
    return {"author": author}


def joke_node(state: State):
    """
    根据推荐作家风格创作笑话的节点函数

    Args:
        state (State): 当前状态对象，包含已推荐的作家名称

    Returns:
        dict: 包含创作笑话内容的字典
    """
    # 构造提示词，请求用指定作家的风格写笑话
    prompt = f"⽤作家：{state['author']} 的⻛格，写⼀个100字以内的笑话"
    # 调用大语言模型创作笑话
    joke = llm.invoke(prompt)
    return {"joke": joke}


# 创建状态图构建器
# State: 状态数据结构定义类
builder = StateGraph(State)

# 添加图中的节点
# author_node: 负责推荐作家的节点
# joke_node: 负责创作笑话的节点
builder.add_node(author_node)
builder.add_node(joke_node)

# 定义节点间的执行顺序
# START -> author_node: 从开始节点跳转到推荐作家节点
# author_node -> joke_node: 从推荐作家节点跳转到创作笑话节点
# joke_node -> END: 从创作笑话节点跳转到结束节点
builder.add_edge(START, "author_node")
builder.add_edge("author_node", "joke_node")
builder.add_edge("joke_node", END)

# 创建内存检查点保存器
checkpointer = InMemorySaver()
# 编译图并设置检查点保存器
# checkpointer: 用于保存图执行过程中的状态
graph = builder.compile(checkpointer=checkpointer)
# 打印编译后的图对象信息
# print(graph)


# 正常执⾏⼀个图
import uuid

# 配置图执行的参数
# configurable: 可配置项字典
#   thread_id: 线程ID，使用uuid生成唯一标识符
config = {
    "configurable": {
        "thread_id": uuid.uuid4(),
    }
}
# 调用图的invoke方法执行整个工作流
# 参数:
#   {}: 初始状态为空字典
#   config: 配置信息
# 返回值:
#   state: 执行完成后包含最终状态的字典
state = graph.invoke({}, config)
# 打印推荐的作家名称
print(state["author"])
print()
# 打印创作的笑话内容
print(state["joke"])
